class S_AES:
    S_BOX = [['9','4','A','B'],
           ['D','1','8','5'],
           ['6','2','0','3'],
           ['C','E','F','7']]
    S_BOX_Reverse = [['A','5','9','B'],
                     ['1','7','8','F'],
                     ['6','0','2','3'],
                     ['C','4','D','E']]
    RCon = ["80","30"] #轮常数
    Gf16_mul_table = [
        ['0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0'],
        ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F'],
        ['0', '2', '4', '6', '8', 'A', 'C', 'E', '3', '1', '7', '5', 'B', '9', 'F', 'D'],
        ['0', '3', '6', '5', 'C', 'F', 'A', '9', 'B', '8', 'D', 'E', '7', '4', '1', '2'],
        ['0', '4', '8', 'C', '3', '7', 'B', 'F', '6', '2', 'E', 'A', '5', '1', 'D', '9'],
        ['0', '5', 'A', 'F', '7', '2', 'D', '8', 'E', 'B', '4', '1', '9', 'C', '3', '6'],
        ['0', '6', 'C', 'A', 'B', 'D', '7', '1', '5', '3', '9', 'F', 'E', '8', '2', '4'],
        ['0', '7', 'E', '9', 'F', '8', '1', '6', 'D', 'A', '3', '4', '2', '5', 'C', 'B'],
        ['0', '8', '3', 'B', '6', 'E', '5', 'D', 'C', '4', 'F', '7', 'A', '2', '9', '1'],
        ['0', '9', '1', '8', '2', 'B', '3', 'A', '4', 'D', '5', 'C', '6', 'F', '7', 'E'],
        ['0', 'A', '7', 'D', 'E', '4', '9', '3', 'F', '5', '8', '2', '1', 'B', '6', 'C'],
        ['0', 'B', '5', 'E', 'A', '1', 'F', '4', '7', 'C', '2', '9', 'D', '6', '8', '3'],
        ['0', 'C', 'B', '7', '5', '9', 'E', '2', 'A', '6', '1', 'D', 'F', '3', '4', '8'],
        ['0', 'D', '9', '4', '1', 'C', '8', '5', '2', 'F', 'B', '6', '3', 'E', 'A', '7'],
        ['0', 'E', 'F', '1', 'D', '3', '2', 'C', '9', '7', '6', '8', '4', 'A', 'B', '5'],
        ['0', 'F', 'D', '2', '9', '6', '4', 'B', '1', 'E', 'C', '3', '8', '7', '5', 'A']
    ]#伽罗瓦有限域乘法表
    def __init__(self,k):
        self.plainText = ""
        self.cipherText = ""
        self.key = k
        self.currentText = ""
        self.initialVector = ""
        self.round_Keys = []
        self.Generate_Round_Keys(self.key)
    def changePlainText(self, p):
        self.plainText = p
    def changeCipherText(self, c):
        self.cipherText = c

    def Search_Gf_mul(self,index_A,index_B):
        """
        伽罗瓦乘法查表
        :param index_A:str
        :param index_B:str
        :return:
        """
        return self.Gf16_mul_table[int(index_A,16)][int(index_B,16)]
    def Permutation(self, text):
        #列混淆
        t_0 = self.XOR(text[0],self.Search_Gf_mul('4',text[1]))
        t_1 = self.XOR(text[1],self.Search_Gf_mul('4',text[0]))
        t_2 = self.XOR(text[2],self.Search_Gf_mul('4',text[3]))
        t_3 = self.XOR(text[3],self.Search_Gf_mul('4',text[2]))
        return t_0+t_1+t_2+t_3
    def Permutation_Reverse(self, text):
        #逆列混淆
        r0 = self.XOR(self.Search_Gf_mul('9', text[0]), self.Search_Gf_mul('2', text[1]))
        r1 = self.XOR(self.Search_Gf_mul('2', text[0]), self.Search_Gf_mul('9', text[1]))
        r2 = self.XOR(self.Search_Gf_mul('9', text[2]), self.Search_Gf_mul('2', text[3]))
        r3 = self.XOR(self.Search_Gf_mul('2', text[2]), self.Search_Gf_mul('9', text[3]))
        return r0 + r1 + r2 + r3
    def Shift(self,text):
        """
        行移位函数，第二行循环移位，第一行不变
        :param: 4hex
        :return: ans 4hex
        """
        return text[0]+text[3]+text[2]+text[1]
    def Substution(self,char,SBox):
        #半字节替换操作
        text = bin(int(char, 16))[2:].zfill(4)
        L = int(text[0])*2+int(text[1])
        R = int(text[2])*2+int(text[3])
        return SBox[L][R]
    def Decryption(self):
        """
        解密算法
        秘文：四个hex
        最后得到四个hex
        :return:
        """
        tmp = self.XOR(self.cipherText, self.round_Keys[4] + self.round_Keys[5])# 秘钥相加
        ans_tmp_1 = self.Shift(tmp)  # 行位移
        tmp = ""
        for i in range(0, len(ans_tmp_1)):
            tmp += self.Substution(ans_tmp_1[i], self.S_BOX_Reverse)  # 半字节代替
        ans_tmp_2 = self.XOR(tmp,self.round_Keys[2] + self.round_Keys[3])# 密钥相加
        ans_tmp_3 = self.Permutation_Reverse(ans_tmp_2) # 列混淆
        ######第一轮加密完成
        ans_tmp_4 = self.Shift(ans_tmp_3) #逆行移位
        tmp = ""
        for i in range(0, len(ans_tmp_4)):
            tmp += self.Substution(ans_tmp_4[i], self.S_BOX_Reverse)  # 半字节代替
        ans_tmp_5 = self.XOR(tmp,self.round_Keys[0] + self.round_Keys[1])# 秘钥相加
        self.plainText = ans_tmp_5
    def Encryption(self):
        """
        加密算法
        明文：四个hex
        最后得到四个hex
        :return:
        """
        tmp = ""
        text = self.XOR(self.plainText,(self.round_Keys[0]+self.round_Keys[1]))#秘钥相加
        for i in range(0,len(text)):
            tmp += self.Substution(text[i],self.S_BOX)#半字节代替
        ans_tmp_1 = self.Shift(tmp) # 行位移
        ans_tmp_2 = self.Permutation(ans_tmp_1) # 列混淆
        ans_tmp_3 = self.XOR(ans_tmp_2,self.round_Keys[2]+self.round_Keys[3]) #密钥相加
        ########第一轮完成
        tmp = ""
        for i in range(0,len(ans_tmp_3)):
            tmp += self.Substution(ans_tmp_3[i],self.S_BOX) #半字节替代
        ans_tmp_4 = self.Shift(tmp) # 行位移
        ans = self.XOR(ans_tmp_4,self.round_Keys[4]+self.round_Keys[5]) # 密钥相加
        self.cipherText = ans
        #加密完成
    def XOR(self,s1,s2):
        ans = ""
        """
        按位异或
        """
        assert len(s1) == len(s2)

        for i in range(0, len(s1)):
            tmp = format(int(s1[i], 16) ^ int(s2[i], 16), 'X')
            ans += tmp
        return ans

    def Generate_Round_Keys(self, key):
        """
        密钥生成,得到的是6个8bits（2个hex）的轮密钥
        :param key:
        :return:
        """
        self.round_Keys.clear()
        self.round_Keys.append(key[0:2])#w0
        self.round_Keys.append(key[2:4])#w1
        #print(self.round_Keys)
        key_2_tmp = self.Fucntion_Generate(self.round_Keys[1],0)
        key_2 = self.XOR(self.round_Keys[0],key_2_tmp)
        self.round_Keys.append(key_2)#w2

        key_3 = self.XOR(self.round_Keys[2],self.round_Keys[1])
        self.round_Keys.append(key_3)#w3

        key_4_tmp = self.Fucntion_Generate(self.round_Keys[3],1)
        key_4 = self.XOR(key_4_tmp,self.round_Keys[2])
        self.round_Keys.append(key_4)#w4

        key_5 = self.XOR(self.round_Keys[4],self.round_Keys[3])
        self.round_Keys.append(key_5)#w5

        print(self.round_Keys)
    def Fucntion_Generate(self,key_word,round):
        """
        “g”函数，输入是一个byte，输出得到一个byte
        :param key_word（两个hex）:
        :return: ans（两个hex）
        """
        assert len(key_word)==2
        text = self.Substution(key_word[1],self.S_BOX) + self.Substution(key_word[0],self.S_BOX)
        ans = self.XOR(text,self.RCon[round])
        return ans

    def AsciiEncryption(self, plain, pad_char=' '):
        """
        每两个 ASCII 字符作为一个 16-bit 块加密。
        输出：将每个块的 2 字节（self.cipherText 为 4 hex）转换成两个 latin-1 字符并累加到字符串返回。
        返回值：拼接后的字符串（每块 2 个字符，可能包含不可见字符）
        """
        # 健壮性：若长度为奇数，补一个字符（pad_char 默认为空格）
        if len(plain) % 2 == 1:
            plain = plain + pad_char

        current = ""
        for i in range(0, len(plain), 2):
            pair = plain[i:i + 2]  # 两个字符
            # 把两个字符编码成 2 字节（latin-1 保证 1:1 映射）
            block_hex = pair.encode('latin-1').hex().upper().zfill(4)  # 例如 'AB' -> '4142'
            self.changePlainText(block_hex)  # 设置明文块的 4 hex
            self.Encryption()  # 产生 self.cipherText (4 hex)
            # 将密文的 4 hex 转为原始 2 字节，再用 latin-1 映射为两个字符并拼接
            block_bytes = bytes.fromhex(self.cipherText)
            block_chars = block_bytes.decode('latin-1')
            current += block_chars

        # 把拼接结果放入实例以便后续使用（注意：这是一个包含 latin-1 映射字符的字符串）
        self.cipherText = current
        return current

    def AsciiDecryption(self, cipher):
        """
        逆过程：接受由 AsciiEncryption 返回的字符串（每块 2 个 latin-1 字符），
        将每两个字符转为 4 hex，调用 Decryption()，把还原的明文字节累加并以 str 返回。
        """
        # 输入长度应该是 2 的倍数（每块 2 个字符）
        if len(cipher) % 2 != 0:
            raise ValueError("cipher length must be multiple of 2 (each block is 2 chars)")

        plain_accum = ""
        for i in range(0, len(cipher), 2):
            block_chars = cipher[i:i + 2]  # 两个 latin-1 字符
            # **关键：使用 latin-1 编码回到原始 2 字节**（不要用默认 utf-8）
            block_hex = block_chars.encode('latin-1').hex().upper().zfill(4)
            self.changeCipherText(block_hex)  # 设置 cipherText = 4 hex
            self.Decryption()  # 产生 self.plainText (4 hex)
            plain_bytes = bytes.fromhex(self.plainText)
            plain_accum += plain_bytes.decode('latin-1')  # 恢复成两个字符并拼接

        # 更新实例状态（可选）
        self.plainText = plain_accum
        return plain_accum
    def DoubleEncryption(self,key):
        assert len(key) == 8
        key1 = key[0:4]
        key2 = key[4:8]
        self.key = key1
        self.Generate_Round_Keys(self.key)
        self.Encryption()
        tmp = self.cipherText
        self.changePlainText(tmp)
        self.key = key2
        self.Generate_Round_Keys(self.key)
        self.Encryption()
    def DoubleDecryption(self,key):
        assert len(key) == 8
        key1 = key[4:8]
        key2 = key[0:4]
        self.key = key1
        self.Generate_Round_Keys(self.key)
        self.Decryption()
        tmp = self.plainText
        self.changeCipherText(tmp)
        self.key = key2
        self.Generate_Round_Keys(self.key)
        self.Decryption()

    """
    三重加解密：
    E-D-E，2keys模式
    """
    def TrippleDecryption(self,key):
        assert len(key) == 8
        key1 = key[0:4]
        key2 = key[4:8]
        self.key = key1
        self.Generate_Round_Keys(self.key)
        self.Decryption()
        #tmp = self.plainText
        #self.changeCipherText(tmp)
        self.key = key2
        self.Generate_Round_Keys(self.key)
        self.Encryption()
        #tmp = self.plainText
        #self.changeCipherText(tmp)
        self.key = key1
        self.Generate_Round_Keys(self.key)
        self.Decryption()
    def TrippleEncryption(self,key):
        assert len(key) == 8
        key1 = key[0:4]
        key2 = key[4:8]
        self.key = key1
        self.Generate_Round_Keys(self.key)
        self.Encryption()
        #tmp = self.cipherText #E
        #self.changeCipherText(tmp)
        self.key = key2
        self.Generate_Round_Keys(self.key)
        self.Decryption() #D
        #tmp = self.plainText
        #self.changePlainText(tmp)
        self.key = key1
        self.Generate_Round_Keys(self.key)
        self.Encryption()
    def CBC_mode_Encryption(self,plain,key,initial_vector):
        """
        CBC密码分组链
        输入一组长明文，分成n个4位hex，
        :param plain:
        :param key:
        :return:
        """
        if initial_vector == "":
            self.initialVector = self.PRNG()#初始向量生成
        else:
            self.initialVector = initial_vector
        #字符串分割
        plain_texts = []
        cipher_texts = []
        self.key = key
        if len(plain) % 4 != 0 :
            plain = plain.ljust((4-(len(plain)%4)+len(plain)),"0")
        for i in range(0,len(plain),4):
            plain_texts.append(plain[i:i+4])
        for idx in range(0,len(plain_texts)):
            if idx == 0 :
                tmp = self.XOR(self.initialVector,plain_texts[idx])
            else:
                tmp = self.XOR(cipher_texts[idx-1],plain_texts[idx])

            self.changePlainText(tmp)
            self.Encryption()
            cipher_texts.append(self.cipherText)
        ans = ""
        for cipherText in cipher_texts:
            ans += cipherText
        return ans
    def CBC_mode_Decryption(self,cipher,key,initial_vector):
        """
        密码分组链解密：
        :param cipher:
        :param key:
        :param initial_vector:
        :return:
        """
        # 字符串分割
        plain_texts = []
        cipher_texts = []
        self.initialVector = initial_vector
        if len(cipher) % 4 != 0:
            cipher += "0" * (len(cipher) % 4)
        for i in range(0, len(cipher), 4):
            cipher_texts.append(cipher[i:i + 4])
        for idx in range(0,len(cipher_texts)):
            self.changeCipherText(cipher_texts[idx])
            self.key = key
            self.Decryption()
            if idx == 0:
                tmp = self.XOR(self.initialVector, self.plainText)
            else:
                tmp = self.XOR(cipher_texts[idx - 1], self.plainText)

            plain_texts.append(tmp)
        ans = ""
        for plainText in plain_texts:
            ans += plainText
        return ans
    def PRNG(self):
        #伪随机数生成器
        import random
        vector = random.randint(1,65537)
        vector_str = format(vector,'X')
        return vector_str
if __name__ == "__main__":
    saes = S_AES("abcd")
    saes.changePlainText("123a")
    print("明文为："+saes.plainText)
    saes.Encryption()
    print("密文为："+repr(saes.cipherText))
    saes.changePlainText("")
    saes.Decryption()
    print("解密结果为："+saes.plainText)

    print('-'*60)
    s = S_AES("ABCD")
    s.AsciiEncryption("1234ABCD!@#$6789")
    print("明文为：1234ABCD!@#$6789")
    # 显示 hex（每两个字符对应 4 个 hex）
    print("密文（hex）:", s.cipherText.encode('latin-1').hex().upper())
    # 或显示可调试的 repr（保留转义）
    print("密文（repr）:", repr(s.cipherText))
    s.AsciiDecryption(s.cipherText)
    print("解密结果为："+s.plainText)
    print('-' * 60)
    ss = S_AES("1234")
    ss.changePlainText("123a")
    ss.DoubleEncryption("12345678")
    print("明文为：123a")
    print("密文为：" + ss.cipherText)
    ss.changePlainText("")
    ss.DoubleDecryption("12345678")
    print("解密结果为："+ss.plainText)
    print('-' * 60)
    sss = S_AES("1234")
    sss.changePlainText("123B")
    sss.TrippleEncryption("12345678")
    print("明文为：123B")
    print("密文为：" + sss.cipherText)
    sss.changePlainText("")
    sss.TrippleDecryption("12345678")
    print("解密结果为：" + sss.plainText)
    print('-'*60)
    ans = sss.CBC_mode_Encryption("12345678abcda","abcd","1234")
    print(ans)
    print(sss.CBC_mode_Decryption(ans,"abcd","1234"))
